{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from glove import Glove\n",
    "from glove import Corpus\n",
    "from IPython.core.debugger import set_trace\n",
    "import re"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_home = \"../data4bilstm\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = \"webnlg\"\n",
    "emb_dim = 300"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir = os.path.join(data_home, experiment_name)\n",
    "train_data_path = os.path.join(data_dir, \"train_data.json\")\n",
    "valid_data_path = os.path.join(data_dir, \"valid_data.json\")\n",
    "test_data_dir = os.path.join(data_dir, \"test_data\")\n",
    "test_data_path_dict = {}\n",
    "for path, folds, files in os.walk(test_data_dir):\n",
    "    for file_name in files:\n",
    "        file_path = os.path.join(path, file_name)\n",
    "        file_name = re.match(\"(.*?)\\.json\", file_name).group(1)\n",
    "        test_data_path_dict[file_name] = file_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = json.load(open(train_data_path, \"r\", encoding = \"utf-8\"))\n",
    "valid_data = json.load(open(valid_data_path, \"r\", encoding = \"utf-8\"))\n",
    "test_data_dict = {}\n",
    "for file_name, path in test_data_path_dict.items():\n",
    "    test_data_dict[file_name] = json.load(open(path, \"r\", encoding = \"utf-8\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6222"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "all_data = train_data + valid_data\n",
    "for data in list(test_data_dict.values()):\n",
    "    all_data.extend(data)\n",
    "    \n",
    "corpus = [sample[\"text\"].split(\" \") for sample in all_data]\n",
    "len(corpus)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Glove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_glove_emb(corpus, window = 10, emb_dim = 100, learning_rate = 0.05, epochs = 10, thr_workers = 6):\n",
    "    corpus_model = Corpus()\n",
    "    corpus_model.fit(corpus, window = window)\n",
    "    print('Dict size: %s' % len(corpus_model.dictionary))\n",
    "    print('Collocations: %s' % corpus_model.matrix.nnz)\n",
    "    \n",
    "    glove = Glove(no_components = emb_dim, learning_rate = learning_rate)\n",
    "    glove.fit(corpus_model.matrix, \n",
    "              epochs = epochs,\n",
    "              no_threads = thr_workers, \n",
    "              verbose = True)\n",
    "    glove.add_dictionary(corpus_model.dictionary)\n",
    "    return glove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dict size: 4126\n",
      "Collocations: 191141\n",
      "Performing 10 training epochs with 6 threads\n",
      "Epoch 0\n",
      "Epoch 1\n",
      "Epoch 2\n",
      "Epoch 3\n",
      "Epoch 4\n",
      "Epoch 5\n",
      "Epoch 6\n",
      "Epoch 7\n",
      "Epoch 8\n",
      "Epoch 9\n"
     ]
    }
   ],
   "source": [
    "# glove\n",
    "golve = train_glove_emb(corpus, emb_dim = emb_dim)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save\n",
    "save_path = os.path.join(\"../pretrained_word_emb\", \"glove_{}_{}.emb\".format(emb_dim, experiment_name))\n",
    "golve.save(save_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Register', 0.9819727727207307),\n",
       " ('U', 0.9798825776793576),\n",
       " ('location', 0.9765197346769859),\n",
       " ('title', 0.9749676051364277),\n",
       " ('Battle', 0.9718289957248922),\n",
       " ('Republic', 0.9708853824285265),\n",
       " ('publisher', 0.9707570316108278),\n",
       " ('municipality', 0.967908142328721),\n",
       " ('order', 0.9674162826876535)]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "golve.most_similar('university', number = 10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(4126, 300)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "golve.word_vectors.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quick Start\n",
    "\n",
    "# # get similar words\n",
    "# golve.most_similar('Massachusetts', number = 10)\n",
    "\n",
    "# # emb matrix shape\n",
    "# golve.word_vectors.shape\n",
    "\n",
    "# # get id \n",
    "# golve.dictionary['Virginia']\n",
    "\n",
    "# # # 指定词条词向量\n",
    "# # glove.word_vectors[glove.dictionary['university']]\n",
    "\n",
    "# # save\n",
    "# save_path = os.path.join(data_home, \"pretrained_word_embeddings\", \"glove_100.emb\")\n",
    "# glove.save(save_path)\n",
    "\n",
    "# # load\n",
    "# glove = Glove()\n",
    "# glove = glove.load(save_path)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
